{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adversarial Attacks using API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from advsecurenet.attacks.gradient_based import FGSM, LOTS\n",
    "from advsecurenet.shared.types.configs.attack_configs import (\n",
    "    FgsmAttackConfig,\n",
    "    LotsAttackConfig,\n",
    ")\n",
    "from advsecurenet.shared.types.configs.attack_configs.attacker_config import (\n",
    "    AttackerConfig,\n",
    ")\n",
    "from advsecurenet.models.model_factory import ModelFactory\n",
    "from advsecurenet.datasets.dataset_factory import DatasetFactory\n",
    "from advsecurenet.attacks.attacker import Attacker\n",
    "from advsecurenet.dataloader.data_loader_factory import DataLoaderFactory\n",
    "from advsecurenet.shared.types.configs.preprocess_config import (\n",
    "    PreprocessConfig,\n",
    "    PreprocessStep,\n",
    ")\n",
    "from advsecurenet.shared.types.configs.device_config import DeviceConfig\n",
    "from advsecurenet.utils.adversarial_target_generator import AdversarialTargetGenerator\n",
    "from advsecurenet.datasets.targeted_adv_dataset import AdversarialDataset\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model\n",
    "model = ModelFactory.create_model(\n",
    "    model_name=\"resnet18\", num_classes=10, pretrained=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets define the preprocessing configuration we want to use\n",
    "preprocess_config = PreprocessConfig(\n",
    "    steps=[\n",
    "        PreprocessStep(name=\"Resize\", params={\"size\": 32}),\n",
    "        PreprocessStep(name=\"CenterCrop\", params={\"size\": 32}),\n",
    "        PreprocessStep(name=\"ToTensor\"),\n",
    "        PreprocessStep(\n",
    "            name=\"ToDtype\", params={\"dtype\": \"torch.float32\", \"scale\": True}\n",
    "        ),\n",
    "        PreprocessStep(\n",
    "            name=\"Normalize\",\n",
    "            params={\"mean\": [0.485, 0.456, 0.406], \"std\": [0.229, 0.224, 0.225]},\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Define the dataset\n",
    "dataset = DatasetFactory.create_dataset(\n",
    "    dataset_type=\"cifar10\", preprocess_config=preprocess_config, return_loaded=False\n",
    ")\n",
    "test_data = dataset.load_dataset(train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the dataloder\n",
    "dataloader = DataLoaderFactory.create_dataloader(dataset=test_data, batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the device config\n",
    "device = DeviceConfig(processor=\"mps\")\n",
    "\n",
    "# Define the fgsm config\n",
    "fgsm_config = FgsmAttackConfig(\n",
    "    targeted=False,\n",
    "    epsilon=0.1,\n",
    "    device=device,\n",
    ")\n",
    "\n",
    "# Now we can define the attack\n",
    "attack = FGSM(config=fgsm_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Untargeted FGSM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 1. Using the Attacker\n",
    "\n",
    "You can use the `Attacker` to run the attack and return the adversarial images in the end. This helps you not to worry about the attack loop and focus on the attack parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the attacker\n",
    "attacker_config = AttackerConfig(\n",
    "    model=model,\n",
    "    attack=attack,\n",
    "    dataloader=dataloader,\n",
    "    device=device,\n",
    "    return_adversarial_images=True,\n",
    ")\n",
    "\n",
    "attacker = Attacker(config=attacker_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adv_imgs = attacker.execute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2. Manual Iteration\n",
    "If you prefer to manually run the attack, you can loop through the dataloader to generate the adversarial samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for images, labels in tqdm(dataloader, desc=\"Attacking\"):\n",
    "    model = model.to(device.processor)\n",
    "    images, labels = images.to(device.processor), labels.to(device.processor)\n",
    "    adv_imgs = attack.attack(model, images, labels)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling results\n",
    "bening_pred = model(images).argmax(dim=1)\n",
    "adv_pred = model(adv_imgs).argmax(dim=1)\n",
    "\n",
    "cifar10_classes = [\n",
    "    \"airplane\",\n",
    "    \"automobile\",\n",
    "    \"bird\",\n",
    "    \"cat\",\n",
    "    \"deer\",\n",
    "    \"dog\",\n",
    "    \"frog\",\n",
    "    \"horse\",\n",
    "    \"ship\",\n",
    "    \"truck\",\n",
    "]\n",
    "\n",
    "index = -1\n",
    "\n",
    "# find a index that benign prediction is equal to the label but adversarial prediction is not equal to the label\n",
    "for i in range(len(bening_pred)):\n",
    "    index += 1\n",
    "    if bening_pred[index] == labels[index] and adv_pred[index] != labels[index]:\n",
    "        break\n",
    "\n",
    "bening_image = images[index].cpu().numpy().transpose(1, 2, 0)\n",
    "adv_image = adv_imgs[index].cpu().numpy().transpose(1, 2, 0)\n",
    "diff = np.abs(bening_image - adv_image)\n",
    "\n",
    "bening_prediction = cifar10_classes[bening_pred[index]]\n",
    "adv_prediction = cifar10_classes[adv_pred[index]]\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.imshow(bening_image)\n",
    "plt.title(f\"Original Image: {bening_prediction}\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.imshow(adv_image)\n",
    "plt.title(f\"Adversarial Image: {adv_prediction}\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.imshow(diff)\n",
    "plt.title(\"Difference\")\n",
    "plt.axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Targeted FGSM\n",
    "\n",
    "We can also run the targeted FGSM Attack. We will utilize the `AdversarialTargetGenerator` to generate the target labels for the attack."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create adversarial target generator\n",
    "target_generator = AdversarialTargetGenerator()\n",
    "\n",
    "# Adversarial Target Generator uses indices mapping to generate target labels. By setting overwrite=True, it will overwrite the existing target labels if they exist.\n",
    "target_labels = target_generator.generate_target_labels(data=test_data, overwrite=True)\n",
    "\n",
    "# We can create a new dataset with the adversarial target labels. This will be used during the attack.\n",
    "adv_data = AdversarialDataset(\n",
    "    base_dataset=test_data,\n",
    "    target_labels=target_labels,\n",
    ")\n",
    "\n",
    "# Since we have a new dataset, we need to create a new dataloader\n",
    "targeted_dataloader = DataLoaderFactory.create_dataloader(\n",
    "    dataset=adv_data, batch_size=32\n",
    ")\n",
    "\n",
    "targeted_fgsm_config = FgsmAttackConfig(\n",
    "    targeted=True,\n",
    "    epsilon=0.1,\n",
    "    device=device,\n",
    ")\n",
    "\n",
    "targeted_fgsm = FGSM(config=targeted_fgsm_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Option 1. Using the Attacker\n",
    "From this point everything is the same as the untargeted attack. We again have two options to run the attack. Either use the `Attacker` or manually iterate through the dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "targeted_attacker_config = AttackerConfig(\n",
    "    model=model,\n",
    "    attack=targeted_fgsm,\n",
    "    dataloader=targeted_dataloader,\n",
    "    device=device,\n",
    "    return_adversarial_images=True,\n",
    ")\n",
    "\n",
    "targeted_fgsm_attacker = Attacker(config=targeted_attacker_config)\n",
    "\n",
    "# Now we can execute the attack\n",
    "targeted_adv_imgs = targeted_fgsm_attacker.execute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2. Manual Iteration\n",
    "\n",
    "This step is almost identical to the untargeted attack. The only difference is here we need to loop through the new dataloader that contains the target labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for images, labels, target_labels, _ in tqdm(targeted_dataloader, desc=\"Attacking\"):\n",
    "    model = model.to(device.processor)\n",
    "    images, labels, target_labels = (\n",
    "        images.to(device.processor),\n",
    "        labels.to(device.processor),\n",
    "        target_labels.to(device.processor),\n",
    "    )\n",
    "    # Note that we are passing the target labels to the attack function\n",
    "    adv_imgs = attack.attack(model, images, target_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOTS Attack\n",
    "\n",
    "Targeted attacks need target labels to be specified, which is used by the attack to get closer to. `LOTS` attack is a targeted attack, in addition to the target labels it also expects the target images to be specified. With the huge dataset, it is not feasible to specify the target images manually. Therefore, `advsecurenet` provides a way to generate the target images automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to generate target images and labels for the targeted attack\n",
    "target_images, target_labels = target_generator.generate_target_images_and_labels(\n",
    "    data=test_data, overwrite=True\n",
    ")\n",
    "\n",
    "# Again lets create a new dataset with the target images and labels\n",
    "lots_dataset = AdversarialDataset(\n",
    "    base_dataset=test_data,\n",
    "    target_images=target_images,\n",
    "    target_labels=target_labels,\n",
    ")\n",
    "\n",
    "# And a new dataloader\n",
    "lots_dataloader = DataLoaderFactory.create_dataloader(\n",
    "    dataset=lots_dataset, batch_size=32\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 1. Using the Attacker\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOTS attack expects a deep feature layer to be passed as an argument. We can get the deep feature layer from the model\n",
    "model.get_layer_names()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets define the lots config. We pick fc as the deep feature layer\n",
    "lots_config = LotsAttackConfig(\n",
    "    device=device, epsilon=0.1, targeted=True, deep_feature_layer=\"fc\"\n",
    ")\n",
    "\n",
    "# And create the attack\n",
    "lots_attack = LOTS(config=lots_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets define the attacker config again\n",
    "lots_attacker_config = AttackerConfig(\n",
    "    model=model,\n",
    "    attack=lots_attack,\n",
    "    dataloader=lots_dataloader,\n",
    "    device=device,\n",
    "    return_adversarial_images=True,\n",
    ")\n",
    "\n",
    "lots_attacker = Attacker(config=lots_attacker_config)\n",
    "\n",
    "# Now we can execute the attack\n",
    "lots_adv_imgs = lots_attacker.execute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2. Manual Iteration\n",
    "\n",
    "If you prefer to manually run the attack, you can loop through the dataloader to generate the adversarial samples. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for images, labels, target_labels, target_images in tqdm(\n",
    "    lots_dataloader, desc=\"Attacking\"\n",
    "):\n",
    "    model = model.to(device.processor)\n",
    "    images, labels, target_labels, target_images = (\n",
    "        images.to(device.processor),\n",
    "        labels.to(device.processor),\n",
    "        target_labels.to(device.processor),\n",
    "        target_images.to(device.processor),\n",
    "    )\n",
    "    # Note that we are passing the target labels to the attack function\n",
    "    lots_adv_imgs = lots_attack.attack(model, images, target_labels, target_images)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "adv2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
